import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import accuracy_score, roc_auc_score, classification_report, roc_curve, confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.ensemble import RandomForestClassifier


# 加载训练数据和测试数据
train_data = pd.read_csv('../data/train.csv', header=None)
test_data = pd.read_csv('../data/test2.csv', header=None)

print(train_data.head())
print(test_data.head())

print(train_data.info())


# 假设第0列为Attrition，其他列为特征
columns = ['Attrition'] + [f'Feature_{i}' for i in range(1, train_data.shape[1])]
train_data.columns = columns

# 测试集没有标签列
test_columns = [f'Feature_{i}' for i in range(1, test_data.shape[1] + 1)]
test_data.columns = test_columns
test_data.columns = [f'Feature_{i}' for i in range(1, test_data.shape[1] + 1)]
# 打印确认列名是否正确
print("Train data columns:", train_data.columns.tolist())
print("Test data columns:", test_data.columns.tolist())

# 检查是否有缺失值

print(train_data.isnull().sum())

print(test_data.isnull().sum())
# 将 Attrition 列移除以准备特征矩阵
X_train = train_data.drop(['Attrition','Feature_1', 'Feature_2', 'Feature_3','Feature_4', 'Feature_5', 'Feature_6'], axis=1)
y_train = train_data['Attrition']

# 对所有非数值型特征进行编码
def safe_label_encode(param, le):
    pass


for col in X_train.select_dtypes(include='object').columns:
    le = LabelEncoder()
    X_train[col] = le.fit_transform(X_train[col])
    # 对测试集也做相同变换（如果需要预测）
    if col in test_data.columns:
        test_data[col] = safe_label_encode(test_data[col], le)

X_train_final, X_val, y_train_final, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)



# 训练模型
model = RandomForestClassifier(random_state=42)
model.fit(X_train_final, y_train_final)

# 模型评估
y_pred = model.predict(X_val)
y_prob = model.predict_proba(X_val)[:, 1]

accuracy = accuracy_score(y_val, y_pred)
auc_score = roc_auc_score(y_val, y_prob)

print(f"Accuracy: {accuracy:.4f}")
print(f"AUC Score: {auc_score:.4f}")
print("\nClassification Report:")
print(classification_report(y_val, y_pred))

# 计算混淆矩阵
cm = confusion_matrix(y_val, y_pred)
print("混淆矩阵：\n", cm)

# 确保 y_val 是整数类型
y_val = y_val.astype(int)

# 计算 ROC 曲线
fpr, tpr, thresholds = roc_curve(y_val, y_prob, pos_label=1)
# 绘制ROC曲线
fpr, tpr, thresholds = roc_curve(y_val, y_prob)
plt.figure(figsize=(10, 6))
plt.plot(fpr, tpr, label=f'AUC = {auc_score:.2f}')
plt.plot([0, 1], [0, 1], 'k--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver Operating Characteristic (ROC)')
plt.legend(loc="lower right")
plt.show()